import SimpleITK as sitk
import numpy as np
from . import geoutils
import cv2

def draw_2Dsquare(image_array: np.ndarray, x0, hWidth=5, label=1):
    """
    draw a square area with value $label in YX-plane centering $x0
    :param image_array: ZYX ordered image array
    :param x0: the index tuple respected to image_array, so ZYX order
    :param hWidth: half width of the square area, should be odd
    :return:
    """
    y_start = x0[1] - hWidth + 1
    x_start = x0[2] - hWidth + 1
    for i in range(2 * hWidth - 1):
        y = y_start + i
        for j in range(2 * hWidth - 1):
            x = x_start + j
            if y < 0 or y >= image_array.shape[1] or x < 0 or x >= image_array.shape[2]:
                continue
            image_array[x0[0], y, x] = label


def get_label_coord(image_array: np.ndarray, label_set: set) -> dict:
    """
    return the dict that containing all the pixel coordinates with value in $label_set
    :param label_set:
    :return: dict with keys of label value, values of int32 coord list. ZYX ordered, default ascending by z component
    """

    ret = dict({})

    # calculate the ROI range
    target_z, target_y, target_x = [], [], []
    projection_z = image_array.any(axis=(1, 2))
    projection_y = image_array.any(axis=(0, 2))
    projection_x = image_array.any(axis=(0, 1))

    for i in range(len(projection_z)):
        if projection_z[i]:
            target_z.append(i)

    for j in range(len(projection_y)):
        if projection_y[j]:
            target_y.append(j)

    for k in range(len(projection_x)):
        if projection_x[k]:
            target_x.append(k)


    for i in target_z:
        # value: [coord_sum, cnt]
        temp_dict = dict({})
        for j in target_y:
            for k in target_x:
                value = image_array[i, j, k]
                if value in label_set:
                    if value in temp_dict:
                        temp_dict[value][0] += np.asarray([i, j, k])
                        temp_dict[value][1] += 1
                    else:
                        temp_dict[value] = [np.asarray([i, j, k]), 1]
        for k, v in temp_dict.items():
            v[0] = v[0].astype(np.float32)
            coord = v[0] / v[1]
            if k in ret:
                ret[k].append(coord.astype(np.int32))
            else:
                ret[k] = [coord.astype(np.int32)]

        # print('finished one slice, z:', i)
        # print(len(temp_dict))

    return ret


def landmark_synthesis(sparse_path, out_path):
    sparse_img = sitk.ReadImage(sparse_path)
    image_array = sitk.GetArrayFromImage(sparse_img)
    target_label_set = {1, 2}
    hWidth = 5


    # ================= initialization ========================
    # status element is [oIdx, oCoord, dir], oIdx is the origin point's index in coord_list, dir is np.ndarray
    label_status_dict = dict({})
    for label in target_label_set:
        label_status_dict[label] = [-1, np.zeros(3), np.zeros(3)]

    label_dict = get_label_coord(image_array, target_label_set)
    # test
    assert len(label_dict) == len(target_label_set)
    for k, v in label_dict.items():
        print('label', k, 'has pixel counts: ', len(v))

    start_z, end_z = float('inf'), -float('inf')
    for coord_list in label_dict.values():
        start_z = min(coord_list[0][0], start_z)
        end_z = max(coord_list[-1][0], end_z)

    # print(label_dict)
    # print(start_z, end_z)

    syn_array = np.zeros(image_array.shape)
    # ========================================================


    # should always start at a tagged pixel
    for z in range(start_z, end_z + 1):
        for label in target_label_set:
            coord_list = label_dict[label]
            status = label_status_dict[label]
            curOIdx = status[0]
            curOCoord = status[1]
            curDir = status[2]

            # if current slice has the label, then update the status
            if curOIdx + 1 < len(coord_list) and z == coord_list[curOIdx + 1][0]:
                curOIdx += 1
                curOCoord = coord_list[curOIdx]
                if curOIdx + 1 < len(coord_list):
                    dir = np.asarray(coord_list[curOIdx + 1]) - np.asarray(curOCoord)
                    dir = dir.astype(np.float32)
                    dir /= dir[0]
                    curDir = dir
                label_status_dict[label] = [curOIdx, curOCoord, curDir]

            if curOIdx == -1:
                continue

            # synthesis the label by linear interpolating
            x = curOCoord + curDir * (z - curOCoord[0])
            x = x.astype(np.uint32)
            draw_2Dsquare(syn_array, x, hWidth, label)

    syn_img = sitk.GetImageFromArray(syn_array)
    syn_img.CopyInformation(sparse_img)
    sitk.WriteImage(syn_img, out_path)



def circumferential_section_2D(segment: np.ndarray, ROI, lm1, lm2, base_secId, need_middle=True):
    """
    return the sectionId offset based on anteroseptal

    :param segment: 2D segment array, XY order
    :param ROI: region to traverse, ((start_x, start_y), (end_x, end_y))
    :param lm1: coordinate of the landmark between anterior and anteroseptal, XY order
    :param lm2: coordinate of the landmark between inferoseptal and inferior
    :param base_secId: the secId of anterior
    :return: the section array filled with the correct secId
    """

    center = np.asarray(segment.shape) / 2
    a1, b1, c1 = geoutils.get_line_function_2D(center, lm1)
    a2, b2, c2 = geoutils.get_line_function_2D(center, lm2)
    a3, b3, c3 = geoutils.get_line_function_2D(center, (lm1 + lm2) / 2)
    start_x, start_y = ROI[0]
    end_x, end_y = ROI[1]

    ret = np.zeros(segment.shape)
    for x in range(start_x, end_x + 1):
        for y in range(start_y, end_y + 1):
            if segment[x, y]:
                l1_pos = a1 * x + b1 * y + c1 > 0
                l2_pos = a2 * x + b2 * y + c2 > 0
                l3_pos = a3 * x + b3 * y + c3 > 0
                if l1_pos and l2_pos:
                    ret[x, y] = base_secId + 0
                elif need_middle and (not l1_pos) and l3_pos:
                    ret[x, y] = base_secId + 1
                elif need_middle and l2_pos and (not l3_pos):
                    ret[x, y] = base_secId + 2
                elif (not need_middle) and (not l1_pos) and l2_pos:
                    ret[x, y] = base_secId + 1
                elif (not l1_pos) and (not l2_pos):
                    if need_middle:
                        ret[x, y] = base_secId + 3
                    else:
                        ret[x, y] = base_secId + 2
                elif need_middle and l1_pos and (not l3_pos):
                    ret[x, y] = base_secId + 4
                elif need_middle and (not l2_pos) and l3_pos:
                    ret[x, y] = base_secId + 5
                elif (not need_middle) and l1_pos and (not l2_pos):
                    ret[x, y] = base_secId + 3
                else:
                    print('undefined section!')
                    exit(-2)

    return ret


def landmark_expand(lm_dict: dict, size: int) -> dict:
    """
    expand the landmark coordinate to fill the range [0, size) along z axis, aka the lm_dict key

    :param lm_dict: dictionary of landmark coordinate, key: z-coord, value: (i, j) coord. all dtype is int32
    :param size:
    :return:
    """
    lm_ldict = list(lm_dict.items())

    start_z, start_coord = lm_ldict[0]
    end_z, end_coord = lm_ldict[-1]
    down_delta = (start_coord - lm_ldict[5][1]) / 5
    up_delta = (end_coord - lm_ldict[-6][1]) / 5

    for z in range(start_z - 1, -1, -1):
        coord = start_coord + (start_z - z) * down_delta
        lm_ldict.insert(0, (z, coord.astype(np.int32)))
    for z in range(end_z + 1, size):
        coord = end_coord + (z - end_z) * up_delta
        lm_ldict.append((z, coord.astype(np.int32)))

    return dict(lm_ldict)


def is_apex(img_array, ignore_len=20):
    contours, _ = cv2.findContours(img_array, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 1:
        return True
    bound_cnt = 0
    for ctr in contours:
        if len(ctr) > ignore_len:
            bound_cnt += 1
    return bound_cnt == 1


def get_section(segment_array: np.ndarray, landmark_array: np.ndarray):
    """
    implements the 17-section operation on based on the segment_array and landmarks.
    landmark is of 3D-array type to fit the manual marking. may consider dictionary if predicted by 2D image?

    :param segment_array:
    :param landmark_array: synthesised 3D landmark
    :return: segment with 17 labels
    """

    projection_z = segment_array.any(axis=(1, 2))
    start_z, end_z = 0, 0
    for z in range(len(projection_z)):
        if projection_z[z]:
            start_z = z
            break
    for z in range(len(projection_z) - 1, -1, -1):
        if projection_z[z]:
            end_z = z
            break

    # swap ZYX array to ZXY order!
    segment_array = segment_array.swapaxes(1, 2)
    landmark_array = landmark_array.swapaxes(1, 2)

    label_dict = get_label_coord(landmark_array, {1, 2})
    lm1_dict = dict({})
    lm2_dict = dict({})
    for coord in label_dict[1]:
        lm1_dict[coord[0]] = coord[1:]
    for coord in label_dict[2]:
        lm2_dict[coord[0]] = coord[1:]

    lm1_dict = landmark_expand(lm1_dict, segment_array.shape[0])
    lm2_dict = landmark_expand(lm2_dict, segment_array.shape[0])

    section_array = np.zeros(segment_array.shape)
    apex_started = False
    for z in range(end_z, start_z - 1, -1):
        if not (z in lm1_dict and z in lm2_dict):
            continue

        percentage = (end_z - z) / (end_z - start_z)

        if apex_started or (percentage >= 2 / 3 and is_apex(segment_array[z])):
            section_array[z] = segment_array[z].astype(np.bool8) * 17
            apex_started = True
            continue

        need_middle = False
        if percentage < 1 / 3:
            baseId = 1
            need_middle = True
        elif percentage >= 1 / 3 and percentage < 2 / 3:
            baseId = 7
            need_middle = True
        else:
            baseId = 13
            need_middle = False

        projection_x = segment_array[z].any(1)
        projection_y = segment_array[z].any(0)

        start_x, end_x = 0, 0
        start_y, end_y = 0, 0
        for y in range(len(projection_y)):
            if projection_y[y]:
                start_y = y
                break
        for y in range(len(projection_y) - 1, -1, -1):
            if projection_y[y]:
                end_y = y
                break
        for x in range(len(projection_x)):
            if projection_x[x]:
                start_x = x
                break
        for x in range(len(projection_x) - 1, -1, -1):
            if projection_x[x]:
                end_x = x
                break

        ROI = ((start_x, start_y), (end_x, end_y))

        section_array[z] = circumferential_section_2D(segment_array[z], ROI, lm1_dict[z], lm2_dict[z], baseId, need_middle)

    # swap ZXY back to ZYX
    return section_array.swapaxes(1, 2)




if __name__ == '__main__':
    # TODO: test section_synthesis
    # in_path = 'test_sparse.nii.gz'
    # out_path = 'test_syn.nii.gz'
    #
    # landmark_synthesis(in_path, out_path)

    # TODO: test circumferential_2D
    # seg_array = np.ones((16, 16), dtype=np.float32)
    # lm1 = np.asarray([3., 3.], dtype=np.float32)
    # lm2 = np.asarray([3., 10.], dtype=np.float32)
    #
    # result = circumferential_section_2D(seg_array, [[0, 0], [15, 15]], lm1, lm2, 0, True)
    # print(result.transpose())

    segment_img = sitk.ReadImage(
        r'C:\study\FFR\myProject\data\SAX_59\xyt\case_00383\segmentation.nii.gz')
    landmark_img = sitk.ReadImage(
        r'C:\study\FFR\myProject\mytest\mydata_section\synthesis_landmark\synthesis_00383.nii.gz')


    # TODO: test get_section
    # segment_img = sitk.ReadImage(r'C:\study\FFR\myProject\mytest\mydata_test\rotate\output_shortAxis\case_00184\segmentation.nii.gz')
    # landmark_img = sitk.ReadImage(r'C:\study\FFR\myProject\mytest\mydata_section\synthesis_landmark\synthesis_00184.nii.gz')
    section_array = get_section(sitk.GetArrayFromImage(segment_img).astype(np.uint8), sitk.GetArrayFromImage(landmark_img))
    section_img = sitk.GetImageFromArray(section_array)
    section_img.CopyInformation(segment_img)
    sitk.WriteImage(section_img, 'test_section_00383.nii.gz')
